import copy
import os.path
import cv2
import numpy as np
import onnxruntime as ort
from glob import glob
from tqdm import tqdm
import os
import cv2
import base64
import json

def bbox2json(bboxes,categorys,imgs,image_path,json_dir):
    assert len(bboxes)==len(categorys)
    assert len(image_path)==len(bboxes)

    for i in range(len(image_path)):
        res={"version": "4.2.10",
                "flags": {},}
        json_path = os.path.join(json_dir, image_path[i].split('/')[-1].split('.png')[0] + '.json')
        with open(json_path,encoding='utf8',mode='w') as f:
            res['imagePath']=os.path.basename( image_path[i])
            bbox=bboxes[i]
            category=categorys[i]
            img=imgs[i]
            h,w,c=img.shape
            imag = cv2.imencode('.jpg', img)[1]
            base64_data = str(base64.b64encode(imag))[2:-1]
            #b64=base64.b64encode(img).decode('utf-8')
            res['imageData']=base64_data
            cur_res=[]
            #print(category)
            for j in range(len(bbox)):
                cur_poly=bbox[j]
                cur_res.append({'label':category[j],'points':[[poly[0],poly[1]] for poly in cur_poly],"shape_type": "polygon","flags": {}})
            res['shapes']=cur_res
            # res['time_labeled']=13141516
            # res['labeled']='true'
            im=cv2.imread(image_path[i])
            h,w,c=im.shape
            res['imageHeight']=h
            res['imageWidth']=w
            res=json.dumps(res)
            f.write(res)

def solo_postprocess(indexs,hmap,wh,offset,proto,thresh=0.3):
    index=indexs[0][0].astype(np.int32)
    n,c,h,w=hmap.shape
    masks=[]
    for idx in index:
        ys=idx//w
        xs=idx%w
        score=hmap[0,0,ys,xs]
        if score<thresh:
            continue
        cur_proto=proto[0,idx,:,:]
        masks.append(cur_proto)
    return masks


def get_max_poly(src_mask, return_mask=False):
    new_mask = np.zeros_like(src_mask)
    contours, hier = cv2.findContours(src_mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    index = None
    max_area = 0
    max_poly = None
    for i in range(len(contours)):
        arclen = cv2.arcLength(contours[i], True)

        epsilon = max(3, int(arclen * 0.005))  # 拟合出的多边形与原轮廓最大距离，可以自己设置，这里根据轮廓周长动态设置
        approx = cv2.approxPolyDP(contours[i], epsilon, True)  # 轮廓的多边形拟合
        area = cv2.contourArea(contours[i])  # 计算面积
        if area >= max_area:
            max_area = area
            index = i
            max_poly = approx
    if index is None:
        return None
    new_mask = cv2.drawContours(new_mask, [contours[index]], -1, 1, -1)
    if return_mask:
        return max_poly, new_mask
    return max_poly

class SoloLabel:
    def __init__(self,onnx_path='',thresh=0.5):
        self.sess= ort.InferenceSession(onnx_path)
        self.thresh=thresh
        self.wh=(256,256)

    def infer_img(self,im):
        im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
        h, w, c = im.shape
        im = cv2.resize(im, self.wh)
        im = im / 255.
        im = (im - 0.5) / 0.5
        im = np.transpose(im, [2, 0, 1])
        im = im[np.newaxis, ...]
        outs = self.sess.run(['topk_inds', 'hmap', 'w_h', 'offset', 'proto_p'], {'input': im.astype(np.float32)})
        masks = solo_postprocess(*outs)
        polys=[]
        for mask in masks:
            mask = cv2.resize(mask, (w,h))
            mask = np.where(mask > 0.5, np.ones_like(mask), np.zeros_like(mask))
            # cv2.imshow('a',mask*255)
            # cv2.waitKey(0)
            poly = get_max_poly(mask)
            if poly is None:
                continue
            poly=poly.reshape(-1,2).astype(np.float)
            polys.append(poly)
        return polys

if __name__=='__main__':
    onnx_path='/home/wsl/项目/ars_train_yun/work_dir/work_dir/solo_grain/model.onnx'
    im_path='/media/wsl/a9f0161f-7971-c843-8c81-c68049a0235a/PublicDataSet/小麦/results'
    files=glob('%s/*g'%im_path)
    model=SoloLabel(onnx_path)
    for file in tqdm(files):
        im=cv2.imread(file)
        polys=model.infer_img(copy.deepcopy(im))
        bbox2json([polys],[['0']*len(polys)],[im],[file],im_path)
